Unstack

沿指定轴将输入张量拆分为多个子张量，输出子张量按内存顺序排列。

\[output_i = input[\text{index along axis}] \quad \text{for each slice along } axis\]
输入：

  • input - 输入张量地址。

  • shape - 输入张量各维度大小数组。

  • ndim - 输入张量维度。

  • axis - 拆分的维度索引。

  • data_size - 单个元素字节大小。

输出：

  • output - 输出张量地址数组，长度等于 shape[axis]。

支持平台：

FT78NE MT7004

备注

  • FT78NE 支持 fp32, fp64, int8, int16, int32, cplx64, cplx128

  • MT7004 支持 fp16, fp32, int16, int32, cplx64

共享存储版本:

void unstack_s(void *input, void *output, int *shape, int ndim, int axis, Uint32 data_size, int core_mask)

C调用示例：

 1//FT78NE示例
 2#include <stdio.h>
 3#include <unstack.h>
 4
 5int main(int argc, char* argv[]) {
 6    float *input = (float *)0xA0000000;
 7
 8    float *out0 = (float *)0xC0000000;
 9    float *out1 = (float *)0xC0100000;
10    float *outputs[2] = {out0, out1};
11
12    int shape[2] = {2, 128};
13    int ndim = 2;
14    int axis = 0;
15    int core_mask = 0xff;
16
17    unstack_s(input, outputs, shape, ndim, axis, sizeof(float), core_mask);
18    return 0;
19}

私有存储版本:

void unstack_p(void *input, void *output, int *shape, int ndim, int axis, Uint32 data_size)

C调用示例：

 1//MT7004示例
 2#include <stdio.h>
 3#include <unstack.h>
 4
 5int main(int argc, char* argv[]) {
 6    half *input = (half *)0x10800000;
 7    half *out0 = (half *)0x10810000;
 8    half *out1 = (half *)0x10820000;
 9    half *outputs[2] = {out0, out1};
10
11    int shape[3] = {1, 2, 64};
12    int ndim = 3;
13    int axis = 1;
14
15    unstack_p(input, outputs, shape, ndim, axis, sizeof(half));
16    return 0;
17}